#!/usr/bin/env python3
import rclpy  # ROS 2 Python客户端库
from rclpy.node import Node  # ROS 2节点基类
from livox_interfaces.msg import CustomMsg  # Livox LiDAR自定义消息类型
from sensor_msgs.msg import PointCloud2, PointField  # ROS标准点云消息
import numpy as np  # 用于高效数值计算
import struct  # 用于二进制数据打包/解包
import open3d as o3d  # 用于3D点云处理和可视化
from threading import Thread, Lock  # 用于多线程编程
import time  # 用于时间测量和延时
import multiprocessing  # 用于并行计算

# 定义PointCloud2的字段（X、Y、Z、反射率）
FIELDS = [
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    PointField(name='intensity', offset=12, datatype=PointField.UINT8, count=1),
]
POINT_STEP = 13  # 每个点的字节数（4+4+4+1）
ROW_STEP = 0  # 每行字节数（动态计算）

# 并行处理点云块的函数（优化性能的关键）
def process_chunk(chunk, distance_threshold, exclude_front, area_half_width, area_half_height):
    """
    处理点云块的过滤逻辑
    
    参数:
        chunk: 点云数据块 (numpy数组)
        distance_threshold: 距离阈值 (米)
        exclude_front: 是否排除正前方区域
        area_half_width: 前方区域半宽度 (米)
        area_half_height: 前方区域半高度 (米)
        
    返回:
        过滤后的点云块
    """
    # 1. 距离过滤：使用平方距离比较（避免开平方运算）
    squared_distances = np.sum(chunk[:, :3] **2, axis=1)  # 计算每个点到原点的平方距离
    distance_mask = squared_distances <= (distance_threshold** 2)  # 距离过滤掩码
    
    # 2. 前方区域过滤：排除正前方宽5米×高5米的区域
    if exclude_front:
        # 正前方区域定义：
        # X > 0（传感器前方）
        # Y 在 [-2.5, 2.5] 之间（宽度5米）
        # Z 在 [-2.5, 2.5] 之间（高度5米）
        in_front_mask = (chunk[:, 0] > 0)  # 点在传感器前方（X轴正方向）
        in_width_mask = (np.abs(chunk[:, 1]) <= area_half_width)  # 在宽度范围内
        in_height_mask = (np.abs(chunk[:, 2]) <= area_half_height)  # 在高度范围内
        
        # 合并为前方区域掩码
        front_area_mask = np.logical_and(np.logical_and(in_front_mask, in_width_mask), in_height_mask)
        
        # 最终掩码：距离符合要求且不在前方区域内
        filter_mask = np.logical_and(distance_mask, ~front_area_mask)
    else:
        filter_mask = distance_mask  # 若不启用前方过滤，仅使用距离掩码
    
    return chunk[filter_mask]  # 应用掩码过滤点云

class LidarProcessor(Node):
    def __init__(self):
        """初始化LiDAR处理器节点"""
        super().__init__('lidar_processor')  # 调用父类构造函数，设置节点名称
        
        # 声明并获取ROS参数：输入点云话题名称,订阅节点的准备
        self.declare_parameter('input_cloud', '/livox/lidar_avia')
        """
        声明一个名为input_cloud的节点参数
        设置默认值为/livox/lidar_avia(这通常是激光雷达数据的话题名称）
        声明参数后，用户可以在启动节点时通过命令行或参数文件修改这个值，而无需修改代码
        """

        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value
        """
        获取input_cloud参数的实际值(可能是默认值，也可能是用户指定的值）
        将参数值转换为字符串类型并赋值给input_topic变量
        后续可以用input_topic变量作为订阅话题的名称
        """

        # 创建订阅者，订阅LiDAR数据，回调函数为self.callback
        self.sub = self.create_subscription(
            CustomMsg, input_topic, self.callback, 10
        )
        
        # 创建发布者，发布处理后的点云数据（ROS标准格式）
        self.pcl_pub = self.create_publisher(
            PointCloud2, '/livox/pointcloud2', 10
        )
        
        # 记录日志，显示订阅和发布的话题名称
        self.get_logger().info(f"订阅: {input_topic}，发布可视化点云到: /livox/pointcloud2")

        # 创建Open3D点云对象和可视化窗口
        self.pcd = o3d.geometry.PointCloud()
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window()
        
        # 设置可视化参数：点大小和背景颜色
        self.vis.get_render_option().point_size = 2.0  # 点的大小
        self.vis.get_render_option().background_color = np.array([0, 0, 0])  # 黑色背景
        
        # 创建锁和标志变量，用于多线程同步
        self.point_cloud_lock = Lock()  # 线程锁，保护点云数据
        self.is_visualizing = True  # 可视化线程运行标志
        self.is_first_frame = True  # 是否是第一帧点云
        
        # 启动单独的线程用于点云可视化（避免阻塞主循环）
        self.visualization_thread = Thread(target=self.visualization_loop)
        self.visualization_thread.daemon = True  # 设置为守护线程，主程序退出时自动结束
        self.visualization_thread.start()
        
        # 点云计数和保存相关参数
        self.frame_count = 0  # 帧计数器
        self.save_path = "/path/to/save/pointclouds/"  # 点云保存路径
        self.save_enabled = False  # 是否启用保存功能
        self.use_fixed_color =  True  # 使用聚类颜色
        
        # 过滤参数
        self.distance_threshold = 50.0  # 只保留50米内的点
        self.exclude_front = True       # 启用前方区域过滤（关键参数）
        
        # 前方区域尺寸参数（宽5米×高5米）
        self.area_width = 0           # 区域宽度（米）
        self.area_height = 0         # 区域高度（米）
        self.area_half_width = self.area_width / 2  # 半宽度（2.5米）
        self.area_half_height = self.area_height / 2  # 半高度（2.5米）
        
        # 聚类参数
        self.cluster_eps = 0.2        # DBSCAN邻域半径（米）
        self.cluster_min_points = 30  # 最小聚类点数
        self.enable_clustering = True  # 是否启用聚类
        self.voxel_size = 0.05         # 降采样体素大小（米）
        
        # 并行处理设置
        self.num_processes = max(1, multiprocessing.cpu_count() - 1)

    def callback(self, msg: CustomMsg):
        """回调函数：处理接收到的点云消息"""
        start_time = time.time()  # 记录处理开始时间
        
        # 将接收到的点云消息转换为NumPy数组
        points_array = np.zeros((len(msg.points), 4), dtype=np.float32)
        for i, p in enumerate(msg.points):
            points_array[i, 0] = p.x  # X坐标
            points_array[i, 1] = p.y  # Y坐标
            points_array[i, 2] = p.z  # Z坐标
            points_array[i, 3] = p.reflectivity  # 反射率值
        
        # 并行过滤点云：将点云分成多个块，并行处理
        chunks = np.array_split(points_array, self.num_processes)
        with multiprocessing.Pool(processes=self.num_processes) as pool:
            # 传递新的区域过滤参数
            results = pool.starmap(process_chunk, [(chunk, self.distance_threshold, 
                                                  self.exclude_front, self.area_half_width,
                                                  self.area_half_height) for chunk in chunks])
        filtered_points = np.concatenate(results)  # 合并所有处理结果
        
        # 如果点太少，不进行聚类（优化性能）
        if len(filtered_points) < 100:
            if self.frame_count % 10 == 0:  # 每10帧记录一次日志
                self.get_logger().info(f"过滤后点太少（{len(filtered_points)}点），跳过聚类")
            colors = np.ones((len(filtered_points), 3))  # 全部设为白色
        else:
            # 创建Open3D点云并降采样
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(filtered_points[:, :3])
            downsampled_pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            
            # 执行聚类
            if self.enable_clustering:
                labels = np.array(downsampled_pcd.cluster_dbscan(
                    eps=self.cluster_eps, 
                    min_points=self.cluster_min_points, 
                    print_progress=False
                ))
                
                # 映射回原始点云
                pcd_tree = o3d.geometry.KDTreeFlann(downsampled_pcd)
                full_labels = np.ones(len(filtered_points), dtype=np.int32) * -1
                
                # 对每个原始点，找到最近的降采样点并复制其聚类标签
                for i in range(len(filtered_points)):
                    [k, idx, _] = pcd_tree.search_knn_vector_3d(filtered_points[i, :3], 1)
                    full_labels[i] = labels[idx[0]]
                
                # 生成颜色
                max_label = full_labels.max()
                self.get_logger().info(f"检测到 {max_label + 1} 个聚类（降采样后点云）")
                
                # 使用HSV色彩空间生成更明显的颜色
                colors = np.zeros((len(filtered_points), 3))
                for label in range(-1, max_label + 1):
                    label_mask = full_labels == label
                    if label == -1:
                        colors[label_mask] = [0, 1, 0]  # 噪声点为黑色
                    else:
                        hue = label / (max_label + 1)
                        rgb = self.hsv_to_rgb(0.3, 0.8, 0.9)  # 鲜艳的颜色
                        colors[label_mask] = rgb
            else:
                # 不使用聚类，基于反射率设置颜色
                intensities = filtered_points[:, 3] / 255.0
                colors = np.zeros((len(filtered_points), 3))
                colors[:, 0] = intensities
        
        # 发布ROS点云消息
        pcl_msg = self.create_pointcloud2_msg(msg.header, filtered_points)
        self.pcl_pub.publish(pcl_msg)
        
        # 更新可视化（加锁确保线程安全）
        with self.point_cloud_lock:
            points_np = filtered_points[:, :3]  # 提取坐标数据
            
            # 首次接收到点云时，检查坐标范围
            if self.is_first_frame:
                x_min, y_min, z_min = np.min(points_np, axis=0)
                x_max, y_max, z_max = np.max(points_np, axis=0)
                self.get_logger().info(f"坐标范围: x=[{x_min:.2f}, {x_max:.2f}], y=[{y_min:.2f}, {y_max:.2f}], z=[{z_min:.2f}, {z_max:.2f}]")
                
                # 自动缩放
                scale = 1.0
                if max(abs(x_max), abs(x_min), abs(y_max), abs(y_min), abs(z_max), abs(z_min)) > 1000:
                    scale = 0.001
                    self.get_logger().info(f"应用缩放因子: {scale}")
                    points_np *= scale
            
            # 更新Open3D点云数据
            self.pcd.points = o3d.utility.Vector3dVector(points_np)
            self.pcd.colors = o3d.utility.Vector3dVector(colors)
            
            # 如果启用保存功能，则保存当前帧点云
            if self.save_enabled:
                self.save_current_frame()
            
            # 标记不再是第一帧
            self.is_first_frame = False
        
        # 记录处理时间
        self.frame_count += 1
        process_time = (time.time() - start_time) * 1000  # 转换为毫秒
        self.get_logger().info(f"帧 {self.frame_count}: 原始{len(msg.points)}点 → 过滤后{len(filtered_points)}点 → 处理时间{process_time:.2f}ms")

    def hsv_to_rgb(self, h, s, v):
        """HSV颜色空间转RGB（用于生成更明显的聚类颜色）"""
        h6 = h * 6
        i = np.floor(h6)
        f = h6 - i
        p = v * (1 - s)
        q = v * (1 - s * f)
        t = v * (1 - s * (1 - f))
        
        i = int(i) % 6
        if i == 0:
            return [v, t, p]
        elif i == 1:
            return [q, v, p]
        elif i == 2:
            return [p, v, t]
        elif i == 3:
            return [p, q, v]
        elif i == 4:
            return [t, p, v]
        else:
            return [v, p, q]

    def create_pointcloud2_msg(self, header, points_array):
        """创建ROS标准PointCloud2消息"""
        pcl_msg = PointCloud2()
        pcl_msg.header = header
        pcl_msg.height = 1
        pcl_msg.width = len(points_array)
        pcl_msg.fields = FIELDS
        pcl_msg.is_bigendian = False
        pcl_msg.point_step = POINT_STEP
        pcl_msg.row_step = POINT_STEP * pcl_msg.width
        pcl_msg.is_dense = False
        
        # 将点数据打包为二进制字节流
        data = []
        for point in points_array:
            x, y, z, intensity = point
            data.append(struct.pack('fffB', x, y, z, int(intensity)))
        pcl_msg.data = b''.join(data)
        
        return pcl_msg

    def visualization_loop(self):
        """可视化线程：持续更新点云显示"""
        is_first_visualization = True
        
        while self.is_visualizing:
            with self.point_cloud_lock:
                if is_first_visualization and len(self.pcd.points) > 0:
                    self.vis.add_geometry(self.pcd)
                    
                    # 设置视图参数
                    ctr = self.vis.get_view_control()
                    ctr.set_front([0, 0, 1])
                    ctr.set_up([0, 1, 0])
                    ctr.set_lookat([0, 0, 0])
                    ctr.set_zoom(0.4)
                    
                    is_first_visualization = False
                elif not is_first_visualization:
                    self.vis.update_geometry(self.pcd)
                    self.vis.poll_events()
                    self.vis.update_renderer()
            
            time.sleep(0.03)  # 控制刷新频率

    def save_current_frame(self):
        """保存当前帧点云到文件"""
        filename = f"{self.save_path}/pointcloud_{self.frame_count:06d}.pcd"
        o3d.io.write_point_cloud(filename, self.pcd)
        self.get_logger().info(f"已保存点云: {filename}")

    def destroy_node(self):
        """节点销毁时的清理工作"""
        self.is_visualizing = False
        self.vis.destroy_window()
        super().destroy_node()

def main(args=None):
    """程序入口点"""
    rclpy.init(args=args)
    node = LidarProcessor()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
